//
// Created by mgepahmge on 2025/10/16.
//

#ifndef MYAMG_MATRIXWRITER_H
#define MYAMG_MATRIXWRITER_H
#include <fstream>
#include <iomanip>
#include <vector>
#include <algorithm>
#include <sstream>
#include <future>
#include <iostream>
#include <unordered_map>
#include <map>

namespace myamg {
    class MtxProcessor {
    public:
        template<typename Integer, typename Real>
        static void writeCSR(const std::string &fileName, const Integer *rowPtr, const Integer *colIndex,
                             const Real *values, const size_t nRows, const size_t nElements) {
            std::ofstream file(fileName);
            file << "%%MatrixMarket matrix coordinate real general\n";
            file << nRows << " " << nRows << " " << nElements << "\n";
            file << std::hexfloat;

            for (size_t i = 0; i < nRows; ++i) {
                for (Integer j = rowPtr[i]; j < rowPtr[i + 1]; ++j) {
                    file << (i + 1) << " " << (colIndex[j] + 1) << " " << values[j] << "\n";
                }
            }
        }

        static void readCSRSize(const std::string &fileName, size_t &nRows, size_t &nElements) {
            std::ifstream file(fileName);
            std::string line;

            while (std::getline(file, line) && line[0] == '%');

            size_t nCols;
            std::istringstream(line) >> nRows >> nCols >> nElements;
        }

        template<typename Integer, typename Real>
        static void readCSR(const std::string &fileName, Integer *&rowPtr, Integer *&colIndex, Real *&values,
                            const size_t nRows, const size_t nElements) {
            std::ifstream file(fileName);
            if (!file.is_open()) {
                throw std::runtime_error("Cannot open file: " + fileName);
            }

            std::string line;

            while (std::getline(file, line) && !line.empty() && line[0] == '%');

            size_t nCols, nNonZeros;
            std::istringstream(line) >> nCols >> nCols >> nNonZeros;

            rowPtr = new Integer[nRows + 1]();
            colIndex = new Integer[nElements];
            values = new Real[nElements];

            std::vector<std::tuple<Integer, Integer, Real> > entries(nElements);
            for (size_t i = 0; i < nElements; ++i) {
                if (!std::getline(file, line)) {
                    delete[] rowPtr;
                    delete[] colIndex;
                    delete[] values;
                    rowPtr = nullptr;
                    colIndex = nullptr;
                    values = nullptr;
                    throw std::runtime_error("Unexpected end of file");
                }

                line.erase(0, line.find_first_not_of(" \t\r\n"));
                line.erase(line.find_last_not_of(" \t\r\n") + 1);

                std::istringstream iss(line);
                Integer row, col;
                std::string valStr;

                if (!(iss >> row >> col >> valStr)) {
                    delete[] rowPtr;
                    delete[] colIndex;
                    delete[] values;
                    rowPtr = nullptr;
                    colIndex = nullptr;
                    values = nullptr;
                    throw std::runtime_error("Failed to parse line " + std::to_string(i) + ": " + line);
                }

                char *endptr;
                Real val = static_cast<Real>(std::strtod(valStr.c_str(), &endptr));

                if (endptr == valStr.c_str()) {
                    delete[] rowPtr;
                    delete[] colIndex;
                    delete[] values;
                    rowPtr = nullptr;
                    colIndex = nullptr;
                    values = nullptr;
                    throw std::runtime_error("Failed to parse value at line " + std::to_string(i) + ": " + valStr);
                }

                entries[i] = {row - 1, col - 1, val};
            }

            std::sort(entries.begin(), entries.end());

            for (size_t i = 0; i < nElements; ++i) {
                rowPtr[std::get<0>(entries[i]) + 1]++;
                colIndex[i] = std::get<1>(entries[i]);
                values[i] = std::get<2>(entries[i]);
            }

            for (size_t i = 1; i <= nRows; ++i) {
                rowPtr[i] += rowPtr[i - 1];
            }

            file.close();
        }

        template<typename Real>
        static void writeVector(const std::string &fileName, const Real *values, const size_t nElements) {
            std::ofstream file(fileName);
            file << "%%MatrixMarket matrix array real general\n";
            file << nElements << " 1\n";
            file << std::setprecision(16) << std::scientific;

            for (size_t i = 0; i < nElements; ++i) {
                file << values[i] << "\n";
            }
        }

        static void readVectorSize(const std::string &fileName, size_t &nElements) {
            std::ifstream file(fileName);
            std::string line;

            while (std::getline(file, line) && line[0] == '%');

            size_t nCols;
            std::istringstream(line) >> nElements >> nCols;
        }

        template<typename Real>
        static void readVector(const std::string &fileName, Real *&values, const size_t nElements) {
            std::ifstream file(fileName);
            if (!file.is_open()) {
                throw std::runtime_error("Cannot open file: " + fileName);
            }

            std::string line;

            while (std::getline(file, line) && !line.empty() && line[0] == '%');

            size_t nRows, nCols;
            std::istringstream(line) >> nRows >> nCols;

            for (size_t i = 0; i < nElements; ++i) {
                if (!std::getline(file, line)) {
                    delete[] values;
                    values = nullptr;
                    throw std::runtime_error("Unexpected end of file");
                }

                line.erase(0, line.find_first_not_of(" \t\r\n"));
                line.erase(line.find_last_not_of(" \t\r\n") + 1);

                char *endptr;
                values[i] = static_cast<Real>(std::strtod(line.c_str(), &endptr));

                if (endptr == line.c_str()) {
                    delete[] values;
                    values = nullptr;
                    throw std::runtime_error("Failed to parse line " + std::to_string(i) + ": " + line);
                }
            }
            file.close();
        }
    };

    template<typename IntegerType, typename FloatType>
    class CsrMatrixReader {
    private:
        size_t nDim;
        size_t nData;
        std::vector<IntegerType> indexPoints;
        std::vector<IntegerType> indices;
        std::vector<FloatType> data;
        std::vector<FloatType> rhs;
        std::vector<FloatType> x;

        std::future<void> indexPointsFuture;
        std::future<void> indicesFuture;
        std::future<void> dataFuture;
        std::future<void> rhsFuture;
        std::future<void> xFuture;

        std::string trim(const std::string &str) {
            size_t start = str.find_first_not_of(" \t\r\n");
            if (start == std::string::npos) return "";
            size_t end = str.find_last_not_of(" \t\r\n");
            return str.substr(start, end - start + 1);
        }

        bool getNextValidLine(std::istringstream &stream, std::string &line) {
            while (std::getline(stream, line)) {
                line = trim(line);
                if (!line.empty()) {
                    return true;
                }
            }
            return false;
        }

        void readIntegerVector(std::istringstream &stream, std::vector<IntegerType> &vec) {
            std::string line;
            while (std::getline(stream, line)) {
                line = trim(line);

                if (line.empty()) {
                    continue;
                }

                if (line[0] == '#') {
                    break;
                }

                std::istringstream iss(line);
                IntegerType value;
                while (iss >> value) {
                    vec.push_back(value);
                }
            }
            vec.shrink_to_fit();
        }

        void readFloatVector(std::istringstream &stream, std::vector<FloatType> &vec) {
            std::string line;
            while (std::getline(stream, line)) {
                line = trim(line);

                if (line.empty()) {
                    continue;
                }

                if (line[0] == '#') {
                    break;
                }

                std::istringstream iss(line);
                FloatType value;
                while (iss >> value) {
                    vec.push_back(value);
                }
            }
            vec.shrink_to_fit();
        }

        std::string extractSection(const std::string &content, const std::string &startMarker) {
            size_t startPos = content.find(startMarker);
            if (startPos == std::string::npos) {
                return "";
            }

            size_t lineEnd = content.find('\n', startPos);
            if (lineEnd == std::string::npos) {
                return "";
            }

            std::string headerLine = content.substr(startPos, lineEnd - startPos);
            size_t numPos = headerLine.find_first_of("0123456789");

            startPos = lineEnd + 1;

            size_t endPos = content.find("\n#", startPos);
            if (endPos == std::string::npos) {
                endPos = content.length();
            }

            return content.substr(startPos, endPos - startPos);
        }

    public:
        explicit CsrMatrixReader(const std::string &filePath) : nDim(0), nData(0) {
            std::ifstream file(filePath);
            if (!file.is_open()) {
                throw std::runtime_error("Failed to open file: " + filePath);
            }

            std::string fileContent((std::istreambuf_iterator<char>(file)),
                                    std::istreambuf_iterator<char>());
            file.close();

            std::istringstream mainStream(fileContent);
            std::string line;

            while (getNextValidLine(mainStream, line)) {
                if (line.find("#Dim") == 0) {
                    size_t pos = line.find_first_of("0123456789");
                    if (pos != std::string::npos) {
                        std::istringstream iss(line.substr(pos));
                        iss >> nDim;
                        rhs.reserve(nDim);
                        x.reserve(nDim);
                        indexPoints.reserve(nDim + 1);
                    }
                } else if (line.find("#Data") == 0) {
                    size_t pos = line.find_first_of("0123456789");
                    if (pos != std::string::npos && pos < line.length()) {
                        std::istringstream iss(line.substr(pos));
                        iss >> nData;
                        data.reserve(nData);
                        indices.reserve(nData);
                    }
                }
            }

            indexPointsFuture = std::async(std::launch::async, [this, fileContent]() {
                std::string section = extractSection(fileContent, "#IndexPoints");
                if (!section.empty()) {
                    std::istringstream stream(section);
                    readIntegerVector(stream, indexPoints);
                }
            });

            indicesFuture = std::async(std::launch::async, [this, fileContent]() {
                std::string section = extractSection(fileContent, "#Indices");
                if (!section.empty()) {
                    std::istringstream stream(section);
                    readIntegerVector(stream, indices);
                }
            });

            dataFuture = std::async(std::launch::async, [this, fileContent]() {
                size_t firstDataPos = fileContent.find("#Data");
                if (firstDataPos == std::string::npos) return;

                size_t lineEnd = fileContent.find('\n', firstDataPos);
                if (lineEnd == std::string::npos) return;

                std::string headerLine = fileContent.substr(firstDataPos, lineEnd - firstDataPos);

                size_t secondDataPos = fileContent.find("#Data", firstDataPos + 5);
                if (secondDataPos != std::string::npos) {
                    std::string section = extractSection(fileContent.substr(secondDataPos), "#Data");
                    if (!section.empty()) {
                        std::istringstream stream(section);
                        readFloatVector(stream, data);
                    }
                }
            });

            rhsFuture = std::async(std::launch::async, [this, fileContent]() {
                std::string section = extractSection(fileContent, "#Rhs");
                if (!section.empty()) {
                    std::istringstream stream(section);
                    readFloatVector(stream, rhs);
                }
            });

            xFuture = std::async(std::launch::async, [this, fileContent]() {
                std::string section = extractSection(fileContent, "#X");
                if (!section.empty()) {
                    std::istringstream stream(section);
                    readFloatVector(stream, x);
                }
            });

            std::cout << "Debug: Async loading initiated" << std::endl;
        }

        size_t getNDim() const { return nDim; }
        size_t getNData() const { return nData; }

        const std::vector<IntegerType> &getIndexPoints() {
            if (indexPointsFuture.valid()) {
                indexPointsFuture.wait();
            }
            return indexPoints;
        }

        const std::vector<IntegerType> &getIndices() {
            if (indicesFuture.valid()) {
                indicesFuture.wait();
            }
            return indices;
        }

        const std::vector<FloatType> &getData() {
            if (dataFuture.valid()) {
                dataFuture.wait();
            }
            std::cout << "Debug: Data vector size after reading: " << data.size() << std::endl;
            return data;
        }

        const std::vector<FloatType> &getRhs() {
            if (rhsFuture.valid()) {
                rhsFuture.wait();
            }
            return rhs;
        }

        const std::vector<FloatType> &getX() {
            if (xFuture.valid()) {
                xFuture.wait();
            }
            return x;
        }

        void printInfo() {
            if (indexPointsFuture.valid()) indexPointsFuture.wait();
            if (indicesFuture.valid()) indicesFuture.wait();
            if (dataFuture.valid()) dataFuture.wait();
            if (rhsFuture.valid()) rhsFuture.wait();
            if (xFuture.valid()) xFuture.wait();

            std::cout << "Dimension: " << nDim << std::endl;
            std::cout << "Data count: " << nData << std::endl;
            std::cout << "IndexPoints size: " << indexPoints.size() << std::endl;
            std::cout << "Indices size: " << indices.size() << std::endl;
            std::cout << "Data size: " << data.size() << std::endl;
            std::cout << "Rhs size: " << rhs.size() << std::endl;
            std::cout << "X size: " << x.size() << std::endl;
        }
    };

    bool writeFeaturesToCSV(const std::vector<std::vector<std::pair<std::string, std::string> > > &features,
                            const std::string &filename);
}

#endif //MYAMG_MATRIXWRITER_H
